package com.raso.example;

import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.model.StreamingResponseHandler;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.model.ollama.OllamaStreamingChatModel;
import dev.langchain4j.model.openai.OpenAiStreamingChatModel;

import java.nio.charset.StandardCharsets;
import java.util.Arrays;
import java.util.concurrent.TimeUnit;

/**
 * 流式响应测试入口
 *
 * @author raos
 * @date 2025-03-17
 */
public class StreamOutputTest {

    public static void main(String[] args) {
        //openaiStreamOutput();
        ollamaStreamOutput();
    }

    public static void ollamaStreamOutput() {
        StreamingChatLanguageModel model = OllamaStreamingChatModel.builder()
                .baseUrl("http://localhost:11434")
                .modelName("deepseek-r1:1.5b")
                .build();

        model.generate("你是谁", new StreamingResponseHandler<AiMessage>() {
            @Override
            public void onNext(String token) {
                try {
                    System.out.println(token);
                    TimeUnit.MILLISECONDS.sleep(5);
                } catch (InterruptedException e) {
                    throw new RuntimeException(e);
                }

            }

            @Override
            public void onError(Throwable error) {
                System.out.println(error);
            }
        });
    }

    // OpenAi流式响应测试
    public static void openaiStreamOutput() {
        StreamingChatLanguageModel model = OpenAiStreamingChatModel.builder()
                .baseUrl("http://langchain4j.dev/demo/openai/v1")
                .apiKey("demo")
                .modelName("gpt-4o-mini")
                .build();

        model.generate("who are you", new StreamingResponseHandler<AiMessage>() {
            @Override
            public void onNext(String token) {
                try {
                    System.out.println(token);
                    TimeUnit.MILLISECONDS.sleep(5);
                } catch (InterruptedException e) {
                    throw new RuntimeException(e);
                }

            }

            @Override
            public void onError(Throwable error) {
                System.out.println(error);
            }
        });
    }

}
